import json
from flask import Flask, jsonify, render_template

app = Flask(__name__)

data_cache = {}


def remap_element_category(element):
    """
    根据新的10分类法，重新映射元素的分类。
    此函数在新数据结构下依然可以正常工作。
    """
    original_cat = element.get('category', 'unknown-category').lower()

    if 'predicted to be ' in original_cat:
        original_cat = original_cat.split('predicted to be ')[-1].strip()

    halogens = {9, 17, 35, 53, 85, 117}
    if element['number'] in halogens:
        return 'halogen'

    mapping = {
        'alkali metal': 'alkali-metal',
        'alkaline earth metal': 'alkaline-earth-metal',
        'lanthanide': 'lanthanide',
        'actinide': 'actinide',
        'transition metal': 'transition-metal',
        'post-transition metal': 'other-metal',
        'metalloid': 'metalloid',
        'polyatomic nonmetal': 'other-nonmetal',
        'diatomic nonmetal': 'other-nonmetal',
        'noble gas': 'noble-gas',
    }

    return mapping.get(original_cat, 'unknown-category')

def load_data():
    """
    在应用启动时加载一次数据并进行分类重映射。
    已适配新的'periodic_table_enriched.json'文件。
    """
    try:
        # --- 核心修改 ---
        # 更新为新爬虫生成的数据文件
        with open("periodic_table_data/periodic_table_enriched.json", "r", encoding="utf-8") as f:
            all_data = json.load(f)

            processed_elements = []
            for elem in all_data['elements']:
                elem['category_slug'] = remap_element_category(elem)
                processed_elements.append(elem)

            data_cache['elements'] = processed_elements
        print("✅ Enriched element data loaded and re-categorized successfully.")
    except FileNotFoundError:
        print("❌ FATAL ERROR: periodic_table_enriched.json not found. Please run the advanced_downloader.py first.")
        data_cache['elements'] = []

# API 端点和路由保持不变，它们的功能不受影响
@app.route("/")
def index():
    return render_template("index.html")

@app.route("/api/elements")
def get_all_elements():
    """API端点：返回经过重分类的元素数据列表"""
    if not data_cache.get('elements'):
        return jsonify({"error": "Element data not loaded on server."}), 500
    return jsonify(data_cache['elements'])

@app.route("/api/phases/<int:temp_k>")
@app.route("/api/phases/<float:temp_k>")
def get_phases_by_temp(temp_k):
    """此函数在新数据结构下依然可以正常工作"""
    if not data_cache.get('elements'):
        return jsonify({"error": "Element data not loaded on server."}), 500

    temp_k = float(temp_k)
    phases = {}
    for element in data_cache['elements']:
        melt = element.get("melt")
        boil = element.get("boil")
        phase = "unknown"
        if melt is None or boil is None:
            if temp_k == 298.15 and element.get("phase"):
                phase = element["phase"].lower()
            else:
                phase = "unknown"
        else:
            if temp_k < melt:
                phase = "solid"
            elif temp_k < boil:
                phase = "liquid"
            else:
                phase = "gas"
        phases[element['number']] = phase
    return jsonify(phases)

# 在应用上下文中加载数据，保持不变
with app.app_context():
    load_data()

if __name__ == "__main__":
    app.run(debug=True, port=5000)